Hi @qubvel,

I attempted to run your notebook example for binary segmentation on our cluster, which has a low-memory GPU. However, the pytorch_lightning trainer encountered issues in this environment, so I had to refactor the training process using a conventional torch training loop.

I tested the modified script (binary_segmentation_intro.py) on a single CPU with a limited run of 10 epochs. While the training process was considerably slow, the implementation appears to be functioning correctly.

I wanted to share this feedback with you and express my appreciation for your valuable contributions to the field.

Regards;
Amir

Hey @omidvarnia , thanks for contribution! I think you opened a PR for a wrong repo, it's a keras repo, can you please reopen against pytorch one? Thanks

Sure, will do it shortly.

Thanks @qubvel! I will address your comments and get back to you.